model_name ='/mnt/nas/shengjie/huggingface_model_local/Qwen-Image-Edit'

from util_for_argparse import get_cuda_port_args

args = get_cuda_port_args()

'''prompt
换成白色上衣短袖，白色下身短裤
'''

import os
from PIL import Image
import torch

from diffusers import QwenImageEditPipeline

pipeline = QwenImageEditPipeline.from_pretrained(model_name,device_map='balanced',torch_dtype=torch.bfloat16)
print("pipeline loaded")
# pipeline.to(torch.bfloat16)
# pipeline.to("cuda")
pipeline.set_progress_bar_config(disable=None)
image = Image.open("./tmp_tryon.jpg").convert("RGB")
prompt = "换成白色上衣短袖，白色下身短裤"
inputs = {
    "image": image,
    "prompt": prompt,
    "generator": torch.manual_seed(0),
    "true_cfg_scale": 4.0,
    "negative_prompt": " ",
    "num_inference_steps": 50,
}

with torch.inference_mode():
    output = pipeline(**inputs)
    output_image = output.images[0]
    output_image.save("tmp_output_image_edit.png")
    print("image saved at", os.path.abspath("tmp_output_image_edit.png"))
